package com.ai.config;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.redis.RedisVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import redis.clients.jedis.JedisPooled;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

@Configuration
@Slf4j
public class VectorStoreConfig {

    // 原始方案（不推荐）
    // 使用 SimpleVectorStore (todo: 内存型(数据存储在内存中，可用save方法保存磁盘), 适用于开发和测试)
    //@Bean
    //public VectorStore vectorStore(EmbeddingModel embeddingModel) {
    //    SimpleVectorStore simpleVectorStore = SimpleVectorStore.builder(embeddingModel).build();
    //
    //    // 定义数据存储目录
    //    File dataStoreDir = new File("D:/data/data_store");
    //
    //    if (dataStoreDir.exists() && dataStoreDir.isDirectory()) {
    //        // 查找所有以 simple_ 开头并以 .json 结尾的文件
    //        File[] vectorFiles = dataStoreDir.listFiles((dir, name) -> name.startsWith("simple_") && name.endsWith(".json"));
    //
    //        if (vectorFiles != null && vectorFiles.length > 0) {
    //            log.info("找到 {} 个向量数据文件", vectorFiles.length);
    //
    //            // 加载所有向量文件
    //            for (File file : vectorFiles) {
    //                try {
    //                    if (file.exists() && file.length() > 0) {
    //                        simpleVectorStore.load(file);
    //                        log.info("成功加载向量数据文件: {}", file.getAbsolutePath());
    //                    }
    //                } catch (Exception e) {
    //                    log.error("加载向量数据文件失败: {}", file.getAbsolutePath(), e);
    //                }
    //            }
    //        } else {
    //            log.warn("在目录 {} 中未找到向量数据文件", dataStoreDir.getAbsolutePath());
    //        }
    //    } else {
    //        log.warn("向量数据存储目录不存在: {}", dataStoreDir.getAbsolutePath());
    //    }
    //
    //    return simpleVectorStore;
    //}

    // 使用自定义的聚合VectorStore实现
    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {

        // 简单的 基于内存，磁盘持久化
        return new AggregatedVectorStore("D:/data/data_store", embeddingModel);

        // 基于redis 向量存储
        //JedisPooled jedisPooled = new JedisPooled("127.0.0.1", 6379);
        //RedisVectorStore redisVectorStore = RedisVectorStore.builder(jedisPooled ,embeddingModel).build();
        //
        //return redisVectorStore;
    }

    /**
     * @param dataStorePath  数据存储目录
     * @param embeddingModel 嵌入模型
     */ // 聚合多个SimpleVectorStore的实现，支持分布式存储同时保持统一检索
    record AggregatedVectorStore(String dataStorePath, EmbeddingModel embeddingModel) implements VectorStore {

        @Override
        public void add(@NotNull List<Document> documents) {
            // 添加文档到最新的向量存储文件中
            SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
            File file = new File(StrUtil.format("{}/simple_{}_{}.json",
                    dataStorePath, embeddingModel.getClass().getSimpleName(), System.currentTimeMillis()));
            FileUtil.mkParentDirs(file);

            if (file.exists() && file.length() > 0) {
                vectorStore.load(file);
            }

            vectorStore.add( documents);
            vectorStore.save(file);
        }

        @Override
        public void delete(@NotNull List<String> idList) {

        }

        @Override
        public void delete(@NotNull Filter.Expression filterExpression) {

        }

        @Override
        public List<Document> similaritySearch(@NotNull SearchRequest request) {
            // 从所有向量存储文件中检索并聚合结果
            List<Document> aggregatedResults = new ArrayList<>();

            File dataStoreDir = new File(dataStorePath);
            if (dataStoreDir.exists() && dataStoreDir.isDirectory()) {
                // 查找所有以 simple_ 开头并以 .json 结尾的文件
                File[] vectorFiles = dataStoreDir.listFiles((dir, name) -> name.startsWith("simple_") && name.endsWith(".json"));

                if (vectorFiles != null && vectorFiles.length > 0) {
                    log.info("找到 {} 个向量数据文件用于相似性检索", vectorFiles.length);

                    // 从每个向量文件中检索并聚合结果
                    for (File file : vectorFiles) {
                        try {
                            if (file.exists() && file.length() > 0) {
                                SimpleVectorStore vectorStore = SimpleVectorStore.builder(embeddingModel).build();
                                vectorStore.load(file);
                                List<Document> fileResults = vectorStore.similaritySearch(request);
                                assert fileResults != null;
                                aggregatedResults.addAll(fileResults);
                                log.debug("从文件 {} 中检索到 {} 个结果", file.getAbsolutePath(), fileResults.size());
                            }
                        } catch (Exception e) {
                            log.error("从向量数据文件检索时失败: {}", file.getAbsolutePath(), e);
                        }
                    }

                    // 按相似度排序并返回topK结果
                    aggregatedResults.sort((d1, d2) -> {
                        double score1 = getDistanceValue(d1);
                        double score2 = getDistanceValue(d2);
                        return Double.compare(score1, score2);
                    });

                    int topK = request.getTopK();
                    if (aggregatedResults.size() > topK) {
                        aggregatedResults = aggregatedResults.subList(0, topK);
                    }
                } else {
                    log.warn("在目录 {} 中未找到向量数据文件", dataStoreDir.getAbsolutePath());
                }
            } else {
                log.warn("向量数据存储目录不存在: {}", dataStoreDir.getAbsolutePath());
            }

            return aggregatedResults;
        }

        /**
         * 安全地从文档元数据中获取距离值
         * @param document 文档
         * @return 距离值，如果不存在则返回Double.MAX_VALUE
         */
        private double getDistanceValue(Document document) {
            Object distanceObj = document.getMetadata().getOrDefault("distance", Double.MAX_VALUE);
            if (distanceObj instanceof Double) {
                return (Double) distanceObj;
            } else if (distanceObj instanceof Float) {
                return (Float) distanceObj;
            } else {
                // 尝试转换为字符串再解析
                try {
                    return Double.parseDouble(distanceObj.toString());
                } catch (NumberFormatException e) {
                    return Double.MAX_VALUE;
                }
            }
        }

    }

    // 或者使用 RedisVectorStore (生产环境)
    // 需确保配置了 spring.data.redis.*
    // @Bean
    // public VectorStore vectorStore(RedisConnectionFactory redisConnectionFactory, EmbeddingModel embeddingModel) {
    //     return new RedisVectorStore(redisConnectionFactory, embeddingModel);
    // }
}
